import numpy as np
import torch
import cv2
from .utils import *
from PIL import Image, ImageDraw



class MaskChamfer:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "mask": ("MASK",),
                "strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
                "iterations": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}),
                "epsilon": ("FLOAT", {"default": 0.01, "min": 0.0, "max": 0.1, "step": 0.001}),
            },
        }

    RETURN_TYPES = ("MASK", "MASK")
    RETURN_NAMES = ("ChamferedMask", "PolygonizedMask")
    FUNCTION = "chamfer_mask"
    CATEGORY = "🌟FJDH/Mask tool"

    def angle_between(self, v1, v2):
        cos_theta = np.clip(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)), -1.0, 1.0)
        angle = np.arccos(cos_theta)
        return angle

    def chamfer_polygon(self, vertices, radius, threshold):
        new_vertices = []
        n = len(vertices)

        for i in range(n):
            A = vertices[i - 1]
            B = vertices[i]
            C = vertices[(i + 1) % n]

            BA = np.array(A) - np.array(B)
            BC = np.array(C) - np.array(B)

            angle = self.angle_between(BA, BC)

            if angle < threshold:
                t1 = B + (BA / np.linalg.norm(BA)) * radius
                t2 = B + (BC / np.linalg.norm(BC)) * radius

                new_vertices.append(t1)
                new_vertices.append(t2)
            else:
                new_vertices.append(B)

        return new_vertices

    def chamfer_mask(self, mask, strength, iterations, epsilon):
        all_chamfered_masks = []
        all_polygonized_masks = []

        for batch_idx, batch_mask in enumerate(mask):
            mask_np = batch_mask.squeeze().cpu().numpy()
            mask_binary = (mask_np > 0).astype(np.uint8)

            contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            polygons = [cv2.approxPolyDP(cnt, epsilon * cv2.arcLength(cnt, True), True) for cnt in contours]

            polygonized_mask = np.zeros_like(mask_binary)
            cv2.fillPoly(polygonized_mask, polygons, 1)
            all_polygonized_masks.append(polygonized_mask)

            threshold = np.pi / 3
            radius = strength * 10

            for iter_idx in range(iterations):
                for poly_idx, polygon in enumerate(polygons):
                    vertices = [point[0] for point in polygon]
                    new_vertices = self.chamfer_polygon(vertices, radius, threshold)
                    new_polygon = np.array(new_vertices, dtype=np.int32).reshape((-1, 1, 2))
                    polygons[poly_idx] = new_polygon

            chamfered_mask = np.zeros_like(mask_binary)
            cv2.fillPoly(chamfered_mask, polygons, 1)

            all_chamfered_masks.append(chamfered_mask)

        chamfered_mask_tensor = np2tensor(all_chamfered_masks)
        polygonized_mask_tensor = np2tensor(all_polygonized_masks)

        return (chamfered_mask_tensor, polygonized_mask_tensor)



class MaskThreshold:
    def __init__(self):
        pass
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "mask": ("MASK",),
                "threshold": ("FLOAT", {
                    "default": 0.5,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.01,
                    "display": "number"
                }),
                "greater_than_threshold": ("BOOLEAN", {"default": True}),
            },
        }
    RETURN_TYPES = ("FLOAT", "BOOLEAN", "MASK")
    RETURN_NAMES = ("Ratio", "Bool", "Mask")
    FUNCTION = "process_mask"
    CATEGORY = "🌟FJDH/Mask tool"
    
    def process_mask(self, mask, threshold, greater_than_threshold):
        mask_area = mask.sum().item()
        total_area = mask.numel()
        ratio = mask_area / total_area

        if greater_than_threshold:
            exceeds_threshold = ratio > threshold
        else:
            exceeds_threshold = ratio < threshold

        if exceeds_threshold:
            empty_mask = torch.zeros_like(mask)
        else:
            empty_mask = mask

        return (ratio, exceeds_threshold, empty_mask)


class BrightnessToMask:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
            },
        }

    RETURN_TYPES = ("MASK", "IMAGE")
    RETURN_NAMES = ("Mask", "BrightnessImage")
    FUNCTION = "convert_brightness_to_mask"
    CATEGORY = "🌟FJDH/Mask tool"

    def convert_brightness_to_mask(self, image):
        if image.ndim == 3:
            image = image.unsqueeze(0)

        brightness_images = []
        masks = []

        for img in image:
            gray_image = cv2.cvtColor(tensor2np(img), cv2.COLOR_RGB2GRAY)
            
            mask_tensor = torch.from_numpy(gray_image) / 255.0

            brightness_image_tensor = torch.stack([torch.from_numpy(gray_image)] * 3, dim=2) 

            masks.append(mask_tensor)
            brightness_images.append(brightness_image_tensor)

        return (torch.cat(masks, dim=0), torch.cat(brightness_images, dim=0))















class BooleanToCombo:
    def __init__(self):
        pass
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "boolean_input": ("BOOLEAN",),
            },
        }
    RETURN_TYPES = (any_type,)
    RETURN_NAMES = ("Combo_output",)
    FUNCTION = "convert_boolean"
    CATEGORY = "🌟FJDH/DWPose tool"

    def convert_boolean(self, boolean_input):
        if boolean_input:
            combo_output = "enable"
        else:
            combo_output = "disable"
        return (combo_output,)

class MaskFilter:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "masks_to_check": ("MASK",),
                "masks_to_add": ("MASK",),
                "threshold": ("FLOAT", {
                    "default": 0.5,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.01,
                    "display": "number"
                }),
                "greater_than_threshold": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("Mask",)
    FUNCTION = "filter_masks"
    CATEGORY = "🌟FJDH/YOLO tool"

    def filter_masks(self, masks_to_check, masks_to_add, threshold, greater_than_threshold):
        filtered_mask = None

        for mask_to_check, mask_to_add in zip(masks_to_check, masks_to_add):
            mask_area = mask_to_check.sum().item()
            total_area = mask_to_check.numel()
            ratio = mask_area / total_area

            if greater_than_threshold:
                exceeds_threshold = ratio > threshold
            else:
                exceeds_threshold = ratio < threshold

            if exceeds_threshold:
                if filtered_mask is None:
                    filtered_mask = mask_to_add
                else:
                    filtered_mask += mask_to_add

        if filtered_mask is None:
            filtered_mask = torch.zeros_like(masks_to_add[0])

        return (filtered_mask,)

class MaskShift:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "mask": ("MASK",),
                "angle": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 360.0, "step": 0.1}),
                "distance": ("INT", {"default": 10, "min": 0, "max": 1000, "step": 1}),
                "fill": ("BOOLEAN", {"default": False}),
            },
        }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("shifted_mask",)
    FUNCTION = "shift_mask"
    CATEGORY = "🌟FJDH/Mask tool"

    def shift_mask(self, mask, angle, distance, fill):
        if mask.dim() == 2:
            mask = torch.unsqueeze(mask, 0)
        
        shifted_masks = []

        for batch_mask in mask:
            batch_mask_np = batch_mask.cpu().numpy()

            shift_x = int(np.cos(np.radians(angle)))
            shift_y = int(np.sin(np.radians(angle)))

            M = np.float32([[1, 0, shift_x], [0, 1, shift_y]])

            if fill:
                accumulated_mask = np.zeros_like(batch_mask_np)
                for _ in range(distance):
                    batch_mask_np = cv2.warpAffine(batch_mask_np, M, (batch_mask_np.shape[1], batch_mask_np.shape[0]), borderMode=cv2.BORDER_REPLICATE)
                    accumulated_mask = np.maximum(accumulated_mask, batch_mask_np)
                batch_mask_np = accumulated_mask
            else:
                shift_x = int(distance * np.cos(np.radians(angle)))
                shift_y = int(distance * np.sin(np.radians(angle)))
                M = np.float32([[1, 0, shift_x], [0, 1, shift_y]])
                batch_mask_np = cv2.warpAffine(batch_mask_np, M, (batch_mask_np.shape[1], batch_mask_np.shape[0]), borderMode=cv2.BORDER_REPLICATE)

            shifted_mask_tensor = torch.from_numpy(batch_mask_np)
            shifted_masks.append(shifted_mask_tensor)

        return (torch.cat(shifted_masks, dim=0),)



class LargestMaskSelector:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "mask": ("MASK",),
            },
        }

    INPUT_IS_LIST = True
    RETURN_TYPES = ("MASK", "MASK")
    RETURN_NAMES = ("Largest_mask", "Merged_mask")
    FUNCTION = "Maxmask"
    CATEGORY = "🌟FJDH/YOLO tool"

    def Maxmask(self, mask):
        max_mask = []
        for item in mask[0]:
            max_mask.append(item.numpy().sum())

        max_index = max_mask.index(max(max_mask))

        largest_mask = mask[0][max_index]

        merged_mask = torch.zeros_like(mask[0][0])
        for i, item in enumerate(mask[0]):
            if i != max_index:
                merged_mask = torch.max(merged_mask, item)

        return largest_mask, merged_mask
    

NODE_CLASS_MAPPINGS = {
    "MaskThreshold": MaskThreshold,
    "BooleanToCombo": BooleanToCombo,
    "MaskFilter": MaskFilter,
    "MaskShift": MaskShift,
    "LargestMaskSelector": LargestMaskSelector,
    "BrightnessToMask": BrightnessToMask,
    "MaskChamfer": MaskChamfer,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "MaskThreshold": "Mask Threshold",
    "BooleanToCombo": "Boolean to Combo",
    "MaskFilter": "Mask Filter",
    "MaskShift": "Mask Shift",
    "LargestMaskSelector": "Largest Mask Selector",
    "BrightnessToMask": "Brightness to Mask",
    "MaskChamfer": "Mask Chamfer",
}